# -*- coding: utf-8 -*-
"""
Created on Mon Dec  4 14:56:04 2023

@author: jetpo
"""

import numpy as np
import matplotlib.pyplot as plt
import os
from PIL import Image, ImageDraw

SAMPLES = 3601

def read_srtm(lat, lon):
    f_name=r'./srtmgl1/N{}E{:0>3}.hgt'.format(lat, lon)
    with open(f_name, 'rb') as hgt_data:
        elevations = np.fromfile(hgt_data, np.dtype('>i2'), SAMPLES * SAMPLES) \
            .reshape((SAMPLES, SAMPLES))
    #lat_range = np.arange(lat, lat + 1/3600 + 1, 1/3600) # 这样是不行的可能出来的数据长度不是3601

    # < hgt文件编排规则 >
    # - 一般的地图, 以地球东北部分为例, 单个瓦片向右经度变大, 向上纬度变大
    # - 但是hgt文件是一个2字节连续编排的纯二进制文件, 如果按顺序读取成二
    #   维数据, 就会变成大的纬度在图片上面:
    #   二维数据左上角--------->经度变大
    #         |                   |
    #         |                   |
    #         ↓                   |
    #   纬度变大(反了)-------二维数据右下角
    # - 因此, 生成纬度数组的时候把大的纬度放在了前面
    #
    # < 单个点读取规则 >
    # - srtm数据, 1个文件跨度为1°（经/纬都是）, 文件内包含3601×3601个点
    # - 每个点由2个字节的无符号整数编码而成
    # - 因此, 右键查看hgt文件时
    
    # < SRTM高程图原始灰度图规则>
    # - 从NASA网站下载到的图片为越黑, 海拔越低
    
    # np.linspace(from, to, count, endpoint)
    # from-to : 数组的起始点
    # count   : 起始点间生成点的总数
    # endpoint: 是否包含终点
    lat_range = np.linspace(lat + 1, lat, 3601, endpoint=True)
    lon_range = np.linspace(lon, lon + 1, 3601, endpoint=True)
    return lat_range, lon_range, elevations


def plot_srtm(ax, lat, lon):
    lat_range, lon_range, elevations = read_srtm(lat, lon)
    
    # np.arange(from, to, res) 
    # from-to: 范围
    # res    : 分辨率
    levels = np.arange(0, 8848, 500)
    
    save_16bit_grayscale(elevations)
    
    # ax.contourf(x, y, z, levels, cmap)
    # x, y  : 横, 纵坐标
    #    z  : 对应的高度值, 总数=x总数×y总数
    # levels: 分辨率规则, 可选, 不选就会自动根据数据范围决定
    # cmap  : 'binary'灰度图渲染风格, 'Spectral':频谱图渲染风格
    #         以灰度显示时越黑越高, 与官方提供图片显示的颜色刚好相反
    #cs = ax.contourf(lon_range, lat_range, elevations, levels, cmap='Spectral') # 频谱图
    cs = ax.contourf(lon_range, lat_range, elevations, 500, cmap='binary') # 灰度图
    return cs

def save_16bit_grayscale(elevations):
    grayImage = Image.new('I;16',(3601, 3601))
    drawGray = ImageDraw.Draw(grayImage)
    for x in range(0, 3600):
        for y in range(0, 3600):
            gray = int(elevations[y, x])
            drawGray.point((x, y), int(gray))
    grayImage.save('e:/test_gray16bit.png')
    print('file saved')          
    
if __name__ == '__main__':
    fig, ax = plt.subplots()
    cs = plot_srtm(ax, 31, 102)
    cbar = fig.colorbar(cs) # 添加colorbar
    plt.show()